import cv2
import numpy as np
import pyrealsense2 as rs

class LaneDetector:
    def __init__(self):
        # RealSense配置
        self.pipeline = rs.pipeline()
        config = rs.config()
        config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
        config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
        self.profile = self.pipeline.start(config)
        
        # 初始化参数
        self.depth_scale = self.profile.get_device().first_depth_sensor().get_depth_scale()
        self.depth_align = rs.align(rs.stream.color)
        self.frame_counter = 0
        self.print_interval = 5  # 每5帧输出一次
        self.prev_left_fit = None
        self.prev_right_fit = None

        # 图像处理参数
        self.gaussian_kernel = (5,5)
        self.canny_thresh = (50, 150)
        self.hough_params = (1, np.pi/180, 20, 30, 25)
        
        # 启用GPU加速
        cv2.ocl.setUseOpenCL(True)

    def _get_roi_mask(self, img):
        """动态调整ROI区域"""
        height, width = img.shape[:2]
        base_top_width = 0.4
        
        # 根据历史斜率调整ROI
        if self.prev_left_fit is not None:
            slope_factor = np.clip(abs(self.prev_left_fit[0])*50, 0.2, 0.4)
            roi_top_width = base_top_width + slope_factor
        else:
            roi_top_width = base_top_width

        vertices = np.array([[
            (int(width*(0.5-roi_top_width/2)), int(height*0.6)),
            (int(width*(0.5+roi_top_width/2)), int(height*0.6)),
            (int(width*0.9), height),
            (int(width*0.1), height)
        ]])
        mask = np.zeros_like(img)
        cv2.fillPoly(mask, vertices, 255)
        return mask

    def _filter_lines(self, lines, width):
        """过滤并分类车道线"""
        left_lines, right_lines = [], []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            if x2 == x1: continue
            slope = (y2-y1)/(x2-x1)
            if abs(slope) < 0.2: continue  # 过滤水平线
            
            # 根据斜率分类
            if slope < 0 and x1 < width/2:
                left_lines.append(line[0])
            elif slope > 0 and x1 > width/2:
                right_lines.append(line[0])
        return left_lines, right_lines

    def _fit_lane(self, lines):
        """带异常处理的车道线拟合"""
        try:
            x = np.concatenate([line[::2] for line in lines])
            y = np.concatenate([line[1::2] for line in lines])
            return np.polyfit(y, x, 1) if len(x)>=2 else None
        except:
            return None

    def _draw_extended_line(self, img, fit, y_start, color):
        """绘制延伸到图像底部的车道线"""
        if fit is not None:
            y_end = img.shape[0]
            x_start = int(fit[0]*y_start + fit[1])
            x_end = int(fit[0]*y_end + fit[1])
            cv2.line(img, (x_start, y_start), (x_end, y_end), color, 3)

    def process_frame(self):
        # 获取对齐的帧数据
        frames = self.pipeline.wait_for_frames()
        aligned_frames = self.depth_align.process(frames)
        color_frame = aligned_frames.get_color_frame()
        depth_frame = aligned_frames.get_depth_frame()
        if not color_frame or not depth_frame: 
            return None

        # 转换为numpy数组
        color_img = np.asanyarray(color_frame.get_data())
        depth_img = np.asanyarray(depth_frame.get_data())

        # 改进的黄色区域检测(LAB空间)
        lab = cv2.cvtColor(color_img, cv2.COLOR_BGR2LAB)
        _, a_channel, _ = cv2.split(lab)
        yellow_mask = cv2.inRange(a_channel, 125, 140)  # 优化后的参数范围

        # 预处理流程
        gray = cv2.cvtColor(color_img, cv2.COLOR_BGR2GRAY)
        masked_gray = cv2.bitwise_and(gray, gray, mask=yellow_mask)
        blurred = cv2.GaussianBlur(masked_gray, self.gaussian_kernel, 0)
        
        # Canny边缘检测(使用GPU加速)
        edges = cv2.UMat(blurred)
        edges = cv2.Canny(edges, *self.canny_thresh)
        roi_edges = cv2.bitwise_and(edges.get(), self._get_roi_mask(edges.get()))

        # 霍夫变换检测
        lines = cv2.HoughLinesP(roi_edges, *self.hough_params)

        # 可视化原始线段
        if lines is not None:
            for line in lines:
                x1, y1, x2, y2 = line[0]
                cv2.line(color_img, (x1,y1), (x2,y2), (255,0,0), 2)

        # 车道线处理
        if lines is not None:
            left_lines, right_lines = self._filter_lines(lines, color_img.shape[1])
            
            # 车道线拟合（使用历史数据降噪）
            left_fit = self._fit_lane(left_lines) if left_lines else self.prev_left_fit
            right_fit = self._fit_lane(right_lines) if right_lines else self.prev_right_fit
            
            if left_fit is not None and right_fit is not None:
                self.prev_left_fit = left_fit
                self.prev_right_fit = right_fit
                y_eval = int(color_img.shape[0] * 0.8)

                # 计算车道位置
                left_x = left_fit[0]*y_eval + left_fit[1]
                right_x = right_fit[0]*y_eval + right_fit[1]
                center_x = (left_x + right_x) / 2

                # 距离计算
                left_dist = (center_x - left_x) * self.depth_scale
                right_dist = (right_x - center_x) * self.depth_scale
                road_width = (right_x - left_x) * self.depth_scale

                # 控制台输出
                if self.frame_counter % self.print_interval == 0:
                    print(f"[Frame {self.frame_counter}] Left: {left_dist:.2f}m | Right: {right_dist:.2f}m | Width: {road_width:.2f}m")
                self.frame_counter += 1

                # 改进的深度测量（区域中值滤波）
                depth_roi = depth_img[y_eval-5:y_eval+5, int(center_x)-5:int(center_x)+5]
                valid_depths = depth_roi[(depth_roi > 0) & (depth_roi < 3/self.depth_scale)]
                if len(valid_depths) > 0:
                    depth_val = np.median(valid_depths) * self.depth_scale
                    intrinsics = self.profile.get_stream(rs.stream.color).as_video_stream_profile().get_intrinsics()
                    center_3d = rs.rs2_deproject_pixel_to_point(intrinsics, [center_x, y_eval], depth_val)

                # 增强可视化
                self._draw_extended_line(color_img, left_fit, y_eval, (0,255,0))
                self._draw_extended_line(color_img, right_fit, y_eval, (0,255,0))
                cv2.circle(color_img, (int(center_x), y_eval), 8, (0,0,255), -1)
                cv2.putText(color_img, f"L:{left_dist:.2f}m | R:{right_dist:.2f}m", (10,30), 
                          cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,0), 2)
                cv2.putText(color_img, f"Road Width: {road_width:.2f}m", (10,60), 
                          cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,0), 2)

        return color_img

    def run(self):
        try:
            while True:
                result = self.process_frame()
                if result is not None:
                    cv2.imshow('Lane Detection', result)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
        finally:
            self.pipeline.stop()
            cv2.destroyAllWindows()

if __name__ == "__main__":
    detector = LaneDetector()
    detector.run()